from fastapi import APIRouter, WebSocket, WebSocketDisconnect
from app.utils.websocket_utils import websocket_manager

router = APIRouter()


@router.websocket("/ws/{user_id}", name="websocket")
async def connect_chat(websocket: WebSocket, user_id: str):
    try:

        await websocket_manager.connect(websocket, user_id)

    except WebSocketDisconnect:
        # 连接断开时移除连接
        del websocket_manager.websocket_connections[user_id]
